package com.raven.play.datalink.jdbc;

import akka.Done;
import akka.NotUsed;
import akka.stream.javadsl.Sink;
import akka.stream.javadsl.Source;
import com.typesafe.config.Config;
import com.typesafe.config.ConfigFactory;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.Statement;
import java.util.concurrent.CompletionStage;

@Getter
@Slf4j
public class JdbcSession {

    private final DataSource dataSource;

    public JdbcSession(DataSource dataSource) {
        this.dataSource = dataSource;
    }

    public static JdbcSession forConfig(String path) {
        Config config = ConfigFactory.load().getConfig(path);
        try {
            DataSource ds = JdbcDataSourceFactory.create(config);
            return new JdbcSession(ds);
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            throw e;
        }
    }

    public <O> Source<O, NotUsed> source(String querySql, ResultSetHandler<O> resultHandler) {
        return Source.fromGraph(new JdbcSourceGraphStage(dataSource, querySql, null, resultHandler){});
    }

    public <O> Source<O, NotUsed> source(String querySql, StatementBuilder<O> statementBuilder, ResultSetHandler<O> resultHandler) {
        return Source.fromGraph(new JdbcSourceGraphStage<O>(dataSource, querySql, statementBuilder, resultHandler){});
    }

    public <I> Sink<I, CompletionStage<Done>> sink(String querySql, StatementBuilder<I> statementBuilder) {
        return Sink.fromGraph(new JdbcSinkGraphStage<I>(dataSource, querySql, statementBuilder, BatchExecutionOptions.defaults()){});
    }

    public <I> Sink<I, CompletionStage<Done>> sink(String querySql, StatementBuilder<I> statementBuilder, BatchExecutionOptions batchExecOpts) {
        return Sink.fromGraph(new JdbcSinkGraphStage<I>(dataSource, querySql, statementBuilder, batchExecOpts){});
    }

    public boolean execute(String sql) {
        try (Connection connection = dataSource.getConnection();
             Statement statement = connection.createStatement()) {
            return statement.execute(sql);
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            throw new RuntimeException(e);
        }
    }
}
